package com.jasmine.limit;

import com.jasmine.core.util.WebUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;

import javax.annotation.PostConstruct;
import javax.servlet.http.HttpServletRequest;
import java.util.ArrayList;
import java.util.List;

/**
 * IP限流,不同接口不共享计数器
 *
 * @author jasmineXz
 */
public class RateLimiterIp {
    private static final Logger log = LoggerFactory.getLogger(RateLimiterIp.class);
    private static final String RATELIMIT_IP_CLASS_PATH = "lua/ratelimit_ip.lua";
    private static final String RATELIMIT_IP_REDIS_KEY = "bjstk_rate_limit_ip:";

    private DefaultRedisScript<Number> redisScript;

    @Autowired
    private RedisTemplate redisTemplate;

    /**
     * 限流
     * @param request request
     * @param rateLimiter 注解
     * @return true:则说明已被限制; false:正常放行
     */
    public boolean limiting (HttpServletRequest request,RateLimiter rateLimiter) {
        List<String> keyList = new ArrayList<>();
        keyList.add(RATELIMIT_IP_REDIS_KEY + request.getRequestURI() + ":" + getIp(request));

        if (redisScript != null) {
            Number num = (Number) redisTemplate.execute(
                    redisScript,
                    keyList,
                    rateLimiter.qps(),
                    rateLimiter.time());

            if (num == null) {
                if (rateLimiter.failFast()) {
                    return false;
                }
            }
            return num.intValue() == 0;
        }
        return false;
    }

    @PostConstruct
    public void init () {
        redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Number.class);
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(RATELIMIT_IP_CLASS_PATH)));
    }

    private String getIp (HttpServletRequest request) {
        return WebUtil.getIP(request);
    }
}
